#include "utils.h"
#include "kpabe_owner.h"
#include "streamforce_crypto_KPAbe.h"
#include "string.h"
#include <stdlib.h>
#include <stdio.h>
#include <iostream>
#include <sstream>
using namespace std;

element_t* maskArray;
//element_t* userDecKeys;

unordered_map<int,dLogMap> owner_discreteLogMap;
unordered_map<int,ptMap> owner_plaintextMap;
unordered_map<int,kpabe_pub_t*> ownerPKMap;
unordered_map<int,kpabe_msk_t*> ownerMKMap;
unordered_map<int, element_wrap> ownerDetKeyMap;
unordered_map<int,int> wsOwnerMap;
unordered_map<int,ptMap> owner_windowKeyMap, owner_randomSumMap;

//streamId -> (columnId -> (windowSize -> element))
unordered_map<int,pptMap> accumulatorMap; 

AES_KEY *aesEncryptionKey, *aesDecryptionKey;

//the ownerId is used to seed the randomness
int initOwner(int streamId, char** attUniverse, int size, int nColumns, int *windowSizes, int nWindows, unsigned char** pubBytes, int*
pubLen){
	kpabe_pub_t* pub;
	kpabe_msk_t* msk;
	pbc_random_set_deterministic(streamId);
	char** atts = parseAttributeUniverse(attUniverse,size);
	kpabe_setup(&pub, &msk, atts);
	toByteArray(kpabe_pub_serialize(pub), pubBytes, pubLen);
	ownerPKMap[streamId] = pub;
	ownerMKMap[streamId] = msk;		

	//init the map
	owner_initDiscreteLogMap(50000, streamId, pub);

	//init the deterministic key X1
	element_wrap tmp;
	element_init_Zr(tmp.element,pub->p);
	element_random(tmp.element);
	ownerDetKeyMap[streamId] = tmp; 
	
	//init beta accumulator for this stream 
	ptMap owkm,am;
	for (int i=0; i<nWindows; i++){
		element_wrap tmp1;
		element_init_Zr(tmp1.element,pub->p);
		element_random(tmp1.element);
		owkm.map[windowSizes[i]] = tmp1;
		for (int j=0; j<nColumns; j++){
			element_wrap tmp;
			element_init_Zr(tmp.element,pub->p);
			element_set(tmp.element, tmp1.element); 		
			accumulatorMap[streamId].map[j].map[windowSizes[i]] = tmp;
		}
	}	
	//accumulatorMap[streamId] = am;
	owner_windowKeyMap[streamId] = owkm;
	
	for (int i=0; i<nColumns; i++){
		element_wrap rs;
		element_init_Zr(rs.element,pub->p);	
		owner_randomSumMap[streamId].map[i] = rs;
	}	
	return 0;
}

//deterministic encryption, using the same X1 for streamId
int enc_det(int streamId, int value, unsigned char **cipher, int *len){
	stringstream out;
	out << value;
	string pt = out.str();
	unsigned char *aesCipher;
	int aesLen;

	kpabe_pub_t* pub = ownerPKMap[streamId];
	encrypt_AES_CBC(pt,&aesCipher, &aesLen, aesEncryptionKey, "0", true);
	element_t testHash;
	element_init_G1(testHash,pub->p);
	element_from_hash(testHash,aesCipher,aesLen);
	
	element_t finalCipher;
	element_init_G1(finalCipher,pub->p);
	element_pow_zn(finalCipher,testHash,ownerDetKeyMap[streamId].element);

	*cipher = (unsigned char*)malloc(element_length_in_bytes(finalCipher)*sizeof(char));
	*len = element_to_bytes(*cipher,finalCipher);
	return 1;  
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_enc_1det
  (JNIEnv *env, jobject obj, jint streamId, jint val){
	unsigned char *cipherBytes;
	int cipherLen;
	enc_det(streamId, val,&cipherBytes,&cipherLen);
	jbyteArray ret = env->NewByteArray(cipherLen);
        env->SetByteArrayRegion(ret, 0, cipherLen, (jbyte*) cipherBytes);
	return ret;	
}

//encryption without masking.
int enc_simple(int streamId, char** atts, int nAtts, int index, int value, unsigned char** cipher, int* len){
	//parse the encryption attributes
	char** encAtts = parseEncryptionAttributes(atts,nAtts);
	if (index%64==63)
	for (int i=0; i<nAtts; i++){
		cout << "att " << i << ": " << encAtts[i] << endl;
	}
	//maping plaintext = gt^value
	element_wrap pt = owner_plaintextMap[streamId].map[value];;
	kpabe_cph_t *cph = kpabe_enc_raw(ownerPKMap[streamId],pt.element,encAtts);
	GByteArray* buf = kpabe_cph_serialize(cph);

	*cipher = buf->data;
	*len = buf->len;
	return 0;
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_enc_1simple
  (JNIEnv *env, jobject obj, jint streamId, jobjectArray encAtts, jint val, jint index){
	int nAtts = env->GetArrayLength(encAtts);
	char** atts = (char**)malloc(nAtts*sizeof(char*));
	for (int i=0; i<nAtts; i++){
		atts[i] =
		(char*)env->GetStringUTFChars((jstring)env->GetObjectArrayElement(encAtts, i), NULL);
	}	
	
	unsigned char* cipherBytes;
	int cipherLen;
	enc_simple(streamId, atts,nAtts, index, val, &cipherBytes, &cipherLen);
	jbyteArray ret = env->NewByteArray(cipherLen);
	env->SetByteArrayRegion(ret, 0, cipherLen, (jbyte*) cipherBytes);
	return ret;
}

//extractin Ep from kpabe_cph_t
JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_extractOriginalCipher
  (JNIEnv *env, jobject obj, jint streamId, jbyteArray cipher){
	env->MonitorEnter(obj);
	int cl = env->GetArrayLength(cipher);
	unsigned char *cipherBytes = (unsigned char*)malloc(cl*sizeof(char));
	env->GetByteArrayRegion(cipher,0,cl,(jbyte*)cipherBytes);
	GByteArray *buf = convertToByteArray(cipherBytes,cl);
	
	//kpabe_cph_t *cph = kpabe_cph_unserialize(ownerPKMap[streamId],buf,1);
	kpabe_cph_t *cph = kpabe_cph_unserialize(cloudPKMap[streamId],buf,1);

	int n = element_length_in_bytes(cph->Ep);
	unsigned char* ep = (unsigned char*)malloc(n*sizeof(char));
	element_to_bytes(ep,cph->Ep);
	jbyteArray ret = env->NewByteArray(n);
	env->SetByteArrayRegion(ret,0,n,(jbyte*)ep);
	env->MonitorExit(obj);
	return ret;
}


int enc_slidingWindow_Aggregate1(int streamId, int columnId, int windowSize, char** atts, int nAtts, int index, int value,  unsigned char** cipher, int* len){
	//parse the encryption attributes
	char** encAtts = parseEncryptionAttributes(atts,nAtts);
	kpabe_pub_t* pub = ownerPKMap[streamId]; 	

	ptMap windowKeys = owner_windowKeyMap[streamId];
	element_t m, m_rm, c_m_rm; 

	if (index%windowSize==0){//reset the accumulator value to the window key

		element_set(accumulatorMap[streamId].map[columnId].map[windowSize].element, windowKeys.map[windowSize].element);
	}
	
	element_t currentSum;
	element_init_Zr(currentSum,pub->p);
	element_set(currentSum,accumulatorMap[streamId].map[columnId].map[windowSize].element);

	//generate a random mask Rm and compute g^{m+Rm}
	element_t rm;
	element_init_Zr(rm,pub->p);
	if (index%windowSize==(windowSize-1))
		element_neg(rm,currentSum);
	else{
		element_random(rm);
	}	
	element_add(accumulatorMap[streamId].map[columnId].map[windowSize].element, currentSum,rm);

	element_init_Zr(m,pub->p);
	element_set_si(m,value);
			
	element_init_Zr(m_rm,pub->p);
	element_add(m_rm,rm,m);
			
	element_init_GT(c_m_rm,pub->p);
	element_pow_zn(c_m_rm,pub->gt,m_rm);
	kpabe_cph_t *cph = kpabe_enc_raw(pub,c_m_rm,encAtts);
	GByteArray* buf = kpabe_cph_serialize(cph);
	*cipher = buf->data;
	*len = buf->len;
			
	return 0;
}


//return two ciphers, encryption of g^{m+r}, g^{sum of r}
//atts1 is for encrypting g^{m+r}, atts2 is for encryptin g^{sum of r}
int enc_slidingWindow_Aggregate3(int streamId, int columnId, char **atts1, int nAtts1, char **atts2, int nAtts2, int index, int value, unsigned char **cipher1, int
*len1, unsigned char **cipher2, int *len2){
	char** encAtts1 = parseEncryptionAttributes(atts1,nAtts1);
	char** encAtts2 = parseEncryptionAttributes(atts2,nAtts2);

	kpabe_pub_t* pub = ownerPKMap[streamId];
	element_t currentSum;
	element_init_Zr(currentSum,pub->p);
	
	element_t rm;
	element_init_Zr(rm,pub->p);
	element_random(rm);
	if (index==0){
		element_init_Zr(owner_randomSumMap[streamId].map[columnId].element,pub->p);
		element_set(owner_randomSumMap[streamId].map[columnId].element,rm);
		element_set(currentSum,rm);
	}
	else{
		element_random(rm);
		element_add(currentSum,owner_randomSumMap[streamId].map[columnId].element,rm);
	}
	
	element_t m,m_rm,g_m_r, g_m_sr;
	element_init_GT(g_m_r,pub->p);	
	element_init_GT(g_m_sr,pub->p);
	element_init_Zr(m,pub->p);
	element_init_Zr(m_rm,pub->p);
	element_set_si(m,value);
	element_add(m_rm,m,rm);

	element_pow_zn(g_m_r,pub->gt,m_rm);
	element_pow_zn(g_m_sr,pub->gt,currentSum);
	element_set(owner_randomSumMap[streamId].map[columnId].element,currentSum);
	
	kpabe_cph_t *cph1 = kpabe_enc_raw(pub,g_m_r,encAtts1);
	kpabe_cph_t *cph2 = kpabe_enc_raw(pub,g_m_sr,encAtts2);
	GByteArray *buf = kpabe_cph_serialize(cph1);
	*cipher1 = buf->data;
	*len1 = buf->len;
	GByteArray *buf1 = kpabe_cph_serialize(cph2);
	*cipher2 = buf1->data;
	*len2 = buf1->len;
	return 0;	
}
JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_enc_1slidingWindow_1Aggregate1
(JNIEnv *env, jobject obj, jint streamId, jint columnId, jint windowSize, jobjectArray encAtts, jint val, jint index){
	int nAtts = env->GetArrayLength(encAtts);
	char** atts = (char**)malloc(nAtts*sizeof(char*));
	for (int i=0; i<nAtts; i++){
		atts[i] =
(char*)env->GetStringUTFChars((jstring)env->GetObjectArrayElement(encAtts, i), NULL);
	}	
	
	unsigned char* cipherBytes;
	int cipherLen;
	enc_slidingWindow_Aggregate1(streamId, columnId, windowSize, atts,nAtts, index, val, &cipherBytes, &cipherLen);
	
	jbyteArray ret = env->NewByteArray(cipherLen);
	env->SetByteArrayRegion(ret, 0, cipherLen, (jbyte*) cipherBytes);
	return ret;
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_enc_1slidingWindow_1Aggregate3
  (JNIEnv *env, jobject obj, jint streamId, jint columnId, jobjectArray atts1, jobjectArray atts2, jint val, jint index, jintArray lens){
	int nAtts1 = env->GetArrayLength(atts1);
	int nAtts2 = env->GetArrayLength(atts2);
	char **encAtts1 = (char**)malloc(nAtts1*sizeof(char*));
	char **encAtts2 = (char**)malloc(nAtts2*sizeof(char*));
	for (int i=0; i<nAtts1; i++)
		encAtts1[i] = (char*)env->GetStringUTFChars((jstring)env->GetObjectArrayElement(atts1, i), NULL);
	for (int i=0; i<nAtts2; i++)
		encAtts2[i] = (char*)env->GetStringUTFChars((jstring)env->GetObjectArrayElement(atts2, i), NULL);

	unsigned char *cipher1, *cipher2;
	int len1, len2;
	enc_slidingWindow_Aggregate3(streamId, columnId, encAtts1, nAtts1, encAtts2, nAtts2,index,val,&cipher1,&len1,&cipher2,&len2);
	jbyteArray ret = env->NewByteArray(len1+len2);
	env->SetByteArrayRegion(ret,0,len1,(jbyte*)cipher1);
	env->SetByteArrayRegion(ret,len1,len2,(jbyte*)cipher2);
	int *tmp = (int*)malloc(2*sizeof(int));
	tmp[0] = len1;
	tmp[1] = len2;
	env->SetIntArrayRegion(lens,0,2,(jint*)tmp);
	return ret; 
}



int generateTransformKey(int streamId, char* policy, unsigned char** key, int userId, int* len){
	char* parsedPolicy = parse_policy_lang(policy);
	kpabe_pub_t* pub = ownerPKMap[streamId];
	//element_init_Zr(decK, pub->p);
	kpabe_prv_t* prvK = kpabe_keygen_proxy(pub,ownerMKMap[streamId],parsedPolicy,userDecKeyMap[userId].element);
	GByteArray* buf = kpabe_prv_serialize(prvK);
	*key = buf->data;
	*len = buf->len;
	return 0;
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_generateTransformKey
  (JNIEnv *env, jobject obj, jint streamId, jstring policy, jint id){
	char* p = (char*)env->GetStringUTFChars(policy, NULL);
	
	unsigned char* keyBytes;
	int keyLen;
	generateTransformKey(streamId, p, &keyBytes, id, &keyLen);
	jbyteArray ret = env->NewByteArray(keyLen);
	env->SetByteArrayRegion(ret,0,keyLen,(jbyte*)keyBytes);

	return ret;	
}

//the owner only has one of such key
int generateUserDetKey(int streamId, unsigned char **key, int *len){
	aesEncryptionKey = get_default_AES_enc_key();
	aesDecryptionKey = get_default_AES_dec_key();
	kpabe_pub_t* pub = ownerPKMap[streamId]; 	

	element_wrap detKey;	
	element_init_Zr(detKey.element,pub->p);
	element_random(detKey.element);
	ownerDetKeyMap[streamId] = detKey; 

	*len = element_length_in_bytes(detKey.element);
	*key = (unsigned char*)malloc((*len)*sizeof(char));
	element_to_bytes(*key,detKey.element);
	return 0;  
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_generateUserDetKey
  (JNIEnv *env, jobject obj, jint streamId){
	unsigned char *keyBytes;
	int keyLen;
	generateUserDetKey(streamId, &keyBytes,&keyLen);
	jbyteArray ret = env->NewByteArray(keyLen);
	env->SetByteArrayRegion(ret,0,keyLen,(jbyte*)keyBytes);
	return ret;
}

int generateUserDecryptionKey(int streamId, unsigned char** key, int* len, int userId){
	element_wrap decK;
	kpabe_pub_t* pub = ownerPKMap[streamId]; 
	element_init_Zr(decK.element,pub->p);
	element_random(decK.element);
	
	*len = element_length_in_bytes(decK.element);
	*key = (unsigned char*)malloc((*len)*sizeof(char));		

	userDecKeyMap[userId] = decK; 
	element_to_bytes(*key, decK.element);
	return 0;
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_generateUserDecryptionKey
  (JNIEnv *env, jobject obj, jint streamId, jint userId){
	unsigned char *keyBytes;
	int keyLen;

	generateUserDecryptionKey(streamId, &keyBytes, &keyLen, userId);
	jbyteArray ret = env->NewByteArray(keyLen);
	env->SetByteArrayRegion(ret,0,keyLen,(jbyte*)keyBytes);
	return ret;
}

int generateUserWindowKey_Agg1(int streamId, int windowSize, unsigned char **key, int *len){
	element_wrap k = owner_windowKeyMap[streamId].map[windowSize];
	*len = element_length_in_bytes(k.element);
	*key = (unsigned char*)malloc((*len)*sizeof(char));
	element_to_bytes(*key,k.element);
	return 0;  
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_generateUserWindowKey_1Aggregate1
  (JNIEnv *env, jobject objc, jint streamId, jint windowSize){
	unsigned char* maskBytes;
	int maskLen;
	generateUserWindowKey_Agg1(streamId, windowSize, &maskBytes, &maskLen);
	jbyteArray ret = env->NewByteArray(maskLen);
	env->SetByteArrayRegion(ret,0,maskLen,(jbyte*)maskBytes);
	return ret;
}

JNIEXPORT jbyteArray JNICALL Java_streamforce_crypto_KPAbe_initOwner
  (JNIEnv *env, jobject obj, jint streamId, jobjectArray atts, jint nColumns, jintArray windowSizes){
	int nAtts = env->GetArrayLength(atts);
	int nWindows = env->GetArrayLength(windowSizes);
	char** attsUniverse = (char**)malloc(nAtts*sizeof(char*));
	bool isCopy = false;
	int *ows = env->GetIntArrayElements(windowSizes,0);
 
	for (int i=0; i<nAtts; i++){
		attsUniverse[i] = (char*)env->GetStringUTFChars((jstring)env->GetObjectArrayElement(atts, i), NULL);
	}		
	unsigned char* pubBytes;
	int pubLen;
	
	initOwner(streamId, attsUniverse,nAtts, nColumns, ows,nWindows, &pubBytes, &pubLen);
	jbyteArray ret = env->NewByteArray(pubLen);
	env->SetByteArrayRegion(ret, 0, pubLen, (jbyte*) pubBytes);
	return ret;
}
